package suffixarray

import (
	"Go-miniSTL/math"
	"Go-miniSTL/mem"
	"unsafe"
)

func num_sort(n, charSize int32, sa, sorta, sort1 []int32) {
	mem.Memset(unsafe.Pointer(&sorta[0]), 0, uintptr(4*charSize))
	for i := int32(0); i < n; i++ {
		sorta[sort1[i]]++
	}
	for i := int32(1); i < charSize; i++ {
		sorta[i] += sorta[i-1]
	}
	for i := int32(0); i < n; i++ {
		sorta[sort1[i]]--
		sa[sorta[sort1[i]]] = i
	}
}

func SA(n, charSize int32, sa, sorta, sort1, sort2 []int32) {
	for k := int32(1); k < n; k <<= 1 {
		rks := int32(-1)
		for i := n - 1; i >= n-k; i-- {
			rks++
			sort2[rks] = i
		}
		for i := int32(0); i < n; i++ {
			if sa[i] >= k {
				rks++
				sort2[rks] = sa[i] - k
			}
		}
		mem.Memset(unsafe.Pointer(&sorta[0]), 0, uintptr(4*charSize))
		for i := int32(0); i < n; i++ {
			sorta[sort1[i]]++
		}
		for i := int32(1); i < charSize; i++ {
			sorta[i] += sorta[i-1]
		}
		for i := n - 1; i >= 0; i-- {
			sorta[sort1[sort2[i]]]--
			sa[sorta[sort1[sort2[i]]]] = sort2[i]
		}
		for i := int32(0); i < n; i++ {
			sort2[i] = sort1[i]
		}
		sort1[sa[0]] = 0
		rks = 1
		for i := int32(1); i < n; i++ {
			sort1[sa[i]] = sort1[sa[i-1]]
			if !(sort2[sa[i]] == sort2[sa[i-1]] && sa[i]+k < n && sa[i-1]+k < n && sort2[sa[i]+k] == sort2[sa[i-1]+k]) {
				sort1[sa[i]]++
				rks++
			}
		}
		if rks == n {
			break
		}
	}
}

func lcp(ss string, n int32, sa, rk, lcps []int32) {
	k := int32(0)
	for i := int32(0); i < n; i++ {
		if k != 0 {
			k--
		}
		j := int32(0)
		if rk[i] > 0 {
			j = sa[rk[i]-1]
		} else {
			lcps[0] = 0
			continue
		}
		for i+k < n && j+k < n && ss[i+k] == ss[j+k] {
			k++
		}
		lcps[rk[i]] = k
	}
}

func New(s string) ([]int32, []int32, []int32) {
	n := int32(len(s))
	charSize := math.Max32(n, 257)
	sorta := make([]int32, charSize)
	sa, rk, lcps := make([]int32, n), make([]int32, n), make([]int32, n)
	sort1, sort2 := make([]int32, n), make([]int32, n)
	for i := int32(0); i < n; i++ {
		sort1[i] = int32(s[i])
	}
	num_sort(n, charSize, sa, sorta, sort1)
	SA(n, charSize, sa, sorta, sort1, sort2)
	for i := int32(0); i < n; i++ {
		rk[sa[i]] = i
	}
	lcp(s, n, sa, rk, lcps)
	return sa, rk, lcps
}
